﻿#include "Decryptor.h"
#include "Eterfree/Platform/Core/Common.h"

#include <exception>
#include <unordered_map>
#include <fstream>
#include <iostream>
#include <filesystem>
#include <atomic>

//#define WIN32_API

#ifdef WIN32_API
//#include <cstring>
#include <sstream>
#else  // !WIN32_API
#include <cstdlib>
#include <cerrno>
#include <cwchar>
#include <clocale>
#include <sstream>
#endif // WIN32_API

#ifdef WIN32_API
#include <Windows.h>

#ifdef UNICODE
using Char = wchar_t;
using String = std::wstring;
using Stream = std::wostringstream;
static auto& cerr = std::wcerr;
#else // !UNICODE
using Char = char;
using String = std::string;
using Stream = std::ostringstream;
static auto& cerr = std::cerr;
#endif // UNICODE

#else // !WIN32_API
//static const char* const SHELL_COMMAND = "C:\\\"Program Files\"\\7-Zip\\7z.exe x \"..\\Workspace\\Eterfree.7z\" -mmt -aoa -bb3";
static const char* const SHELL_COMMAND = "start /min /wait ";
#endif // WIN32_API

static std::wostream& trace(std::wostream& _stream, \
	const wchar_t* _function, const wchar_t* _file, auto _line)
{
	return _stream << L"in " << _function \
		<< L" at " << _file << L':' << _line << L' ';
}

static std::ostream& trace(std::ostream& _stream, \
	const char* _function, const char* _file, auto _line)
{
	return _stream << "in " << _function \
		<< " at " << _file << ':' << _line << ' ';
}

#define TRACEW(stream) trace(stream, __FUNCTIONW__, __FILEW__, __LINE__)
#define TRACEA(stream) trace(stream, __FUNCTION__, __FILE__, __LINE__)

#ifdef UNICODE
#define TRACE TRACEW
#else // !UNICODE
#define TRACE TRACEA
#endif // UNICODE

USING_ETERFREE_SPACE

static bool convert(std::string& _target, \
	const std::string& _source)
{
	_target = _source;
	return true;
}

#ifdef WIN32_API
static bool convert(std::wstring& _target, const std::string& _source)
{
	auto size = static_cast<int>(_source.size());
	if (size <= 0)
	{
		_target.clear();
		return true;
	}

	auto length = ::MultiByteToWideChar(CP_ACP, 0, \
		_source.c_str(), size, NULL, 0);
	if (length != 0)
	{
		_target.assign(length, L'\0');
		::MultiByteToWideChar(CP_ACP, 0, \
			_source.c_str(), size, _target.data(), length);
		return true;
	}

	auto error = ::GetLastError();
	String buffer;
	formatError(buffer, error);

	Stream stream;
	TRACE(stream) << TEXT("MultiByteToWideChar error ") << error \
		<< TEXT(": ") << buffer << std::endl;
	cerr << stream.str();
	return false;
}

static bool convert(std::string& _target, const std::wstring& _source)
{
	auto size = static_cast<int>(_source.size());
	if (size <= 0)
	{
		_target.clear();
		return true;
	}

	auto length = ::WideCharToMultiByte(CP_ACP, 0, \
		_source.c_str(), size, NULL, 0, NULL, NULL);
	if (length != 0)
	{
		_target.assign(length, '\0');
		::WideCharToMultiByte(CP_ACP, 0, _source.c_str(), size, \
			_target.data(), length, NULL, NULL);
		return true;
	}

	auto error = ::GetLastError();
	String buffer;
	formatError(buffer, error);

	Stream stream;
	TRACE(stream) << TEXT("WideCharToMultiByte error ") << error \
		<< TEXT(": ") << buffer << std::endl;
	cerr << stream.str();
	return false;
}

#else // !WIN32_API
static bool convert(std::wstring& _target, const std::string& _source)
{
	auto string = _source.c_str();
	std::mbstate_t state{ 0 };

	auto length = std::mbsrtowcs(nullptr, &string, 0, &state);
	if (length == static_cast<decltype(length)>(-1))
	{
		std::ostringstream stream;
		TRACEA(stream) << "mbsrtowcs error " << errno << std::endl;
		std::cerr << stream.str();
		return false;
	}

	_target.assign(length, L'\0');
	length = std::mbsrtowcs(_target.data(), &string, length, &state);
	if (length == static_cast<decltype(length)>(-1))
	{
		std::ostringstream stream;
		TRACEA(stream) << "mbsrtowcs error " << errno << std::endl;
		std::cerr << stream.str();
		return false;
	}
	return true;
}

static bool convert(std::string& _target, const std::wstring& _source)
{
	auto string = _source.c_str();
	std::mbstate_t state{ 0 };

	auto length = std::wcsrtombs(nullptr, &string, 0, &state);
	if (length == static_cast<decltype(length)>(-1))
	{
		std::ostringstream stream;
		TRACEA(stream) << "wcsrtombs error " << errno << std::endl;
		std::cerr << stream.str();
		return false;
	}

	_target.assign(length, '\0');
	length = std::wcsrtombs(_target.data(), &string, length, &state);
	if (length == static_cast<decltype(length)>(-1))
	{
		std::ostringstream stream;
		TRACEA(stream) << "wcsrtombs error " << errno << std::endl;
		std::cerr << stream.str();
		return false;
	}
	return true;
}
#endif // WIN32_API

static std::string root;
static void initPath()
{
	if (Platform::getImagePath(root))
	{
		std::filesystem::path path = root;
		path.remove_filename();
		if (path.has_parent_path())
			root = path.parent_path().string();
		else
			root.clear();
	}
}

static char table[TABLE_SIZE];
static void initTable() noexcept
{
	for (char index = 0; index < NUMBER_SIZE; ++index)
		table[index] = '0' + index;

	for (char index = 0; index < LETTER_SIZE; ++index)
		table[index + NUMBER_SIZE] = 'A' + index;

	constexpr char OFFSET = NUMBER_SIZE + LETTER_SIZE;
	for (char index = 0; index < LETTER_SIZE; ++index)
		table[index + OFFSET] = 'a' + index;
}

static void initialize() noexcept
{
	try
	{
#ifndef WIN32_API
		std::setlocale(LC_CTYPE, "");
#endif // !WIN32_API

		initPath();
	}
	catch (std::exception&) {}

	initTable();
}

struct Structure
{
	using Config = Decryptor::Config;
	using SizeType = Decryptor::SizeType;
	using IDType = Decryptor::IDType;

	Config _config;
	std::atomic<SizeType> _counter;

	mutable std::mutex _mutex;
	std::unordered_map<IDType, SizeType> _countPool;

	Structure() : _counter(0) {}

	Structure(const Config& _config);

	Structure(Config&& _config) : \
		_config(std::forward<Config>(_config)), \
		_counter(0)
	{
		auto& workspace = this->_config._workspace;
		workspace = root + workspace;
	}

	bool exist() const;

	auto count() const noexcept
	{
		return _counter.load(std::memory_order::relaxed);
	}
	std::optional<SizeType> count(IDType _id) const;
	void count(IDType _id, SizeType _count)
	{
		std::lock_guard lock(_mutex);
		_countPool.insert_or_assign(_id, _count);
	}

	void readCount();
	void writeCount() const
	{
		writeCount(count());
	}
	void writeCount(SizeType _count) const;

#ifdef WIN32_API
	bool getCommand(String& _command) const;
	bool getParameter(String& _parameter, IDType _id);
#else // !WIN32_API
	std::string& getCommand(std::string& _command, IDType _id);
#endif // WIN32_API

	void decrypt(IDType _id);
};

Structure::Structure(const Config& _config) : \
	_config(_config), _counter(0)
{
	auto& workspace = this->_config._workspace;
	workspace.assign(root);
	workspace.append(_config._workspace);
}

bool Structure::exist() const
{
	using namespace std;

	filesystem::path path = _config._workspace;
	path.append(_config._password);
	return filesystem::exists(path);
}

auto Structure::count(IDType _id) const \
-> std::optional<SizeType>
{
	std::lock_guard lock(_mutex);
	if (auto iterator = _countPool.find(_id); \
		iterator != _countPool.end())
		return iterator->second;
	return std::nullopt;
}

void Structure::readCount()
{
	SizeType count = 0;

	auto path = _config._workspace;
	path.append(1, '\\');
	path.append(_config._counter);

	std::ifstream input(path);
	if (input.is_open())
		input >> count;

	_counter.store(count,\
		std::memory_order::relaxed);
}

void Structure::writeCount(SizeType _count) const
{
	auto path = _config._workspace;
	path.append(1, '\\').append(_config._counter);
	std::ofstream output(path);
	if (output.is_open()) output << _count;
}

template <typename _String, typename _SizeType>
static _String& getPassword(_String& _password, _SizeType _count) noexcept
{
	for (auto index = _password.size(); index > 0; \
		--index, _count /= TABLE_SIZE)
		_password[index - 1] = table[_count % TABLE_SIZE];
	return _password;
}

#ifdef WIN32_API
bool Structure::getCommand(String& _command) const
{
	decltype(_config._workspace) command(1, '\"');
	command.append(_config._workspace).append(1, '\\');
	command.append(_config._command).append(1, '\"');
	return convert(_command, command);
}

bool Structure::getParameter(String& _parameter, IDType _id)
{
	auto count = _counter.fetch_add(1, \
		std::memory_order::relaxed);
	this->count(_id, count);

	String workspace;
	if (not convert(workspace, _config._workspace))
		return false;

	_parameter.assign(TEXT(" -o\""));
	_parameter.append(workspace).append(1, TEXT('\\'));

#ifdef UNICODE
	_parameter.append(std::to_wstring(_id)).append(1, L'\"');
#else // !UNICODE
	_parameter.append(std::to_string(_id)).append(1, '\"');
#endif // UNICODE

	String password(_config._length, TEXT('0'));
	getPassword(password, count);
	_parameter.append(TEXT(" -p")).append(password);
	return true;
}

void Structure::decrypt(IDType _id)
{
	String command;
	if (not getCommand(command)) return;

	String parameter;
	if (not getParameter(parameter, _id))
		return;

	//String workspace;
	//if (not convert(workspace, \
	//	_config._workspace)) return;

	//SHELLEXECUTEINFO info;
	//std::memset(&info, 0, sizeof info);
	SHELLEXECUTEINFO info = { 0 };
	info.cbSize = sizeof info;
	info.fMask = SEE_MASK_NOCLOSEPROCESS;
	info.hwnd = ::GetConsoleWindow();
	info.lpVerb = TEXT("open");
	info.lpFile = command.c_str();
	info.lpParameters = parameter.c_str();
	//info.lpDirectory = workspace.c_str();
	info.lpDirectory = NULL;
	info.nShow = SW_HIDE;
	//info.nShow = SW_MINIMIZE;
	info.hInstApp = NULL;

	if (not ShellExecuteEx(&info))
	{
		auto error = ::GetLastError();
		String buffer;
		formatError(buffer, error);

		Stream stream;
		stream << TEXT('[') << _id << TEXT(']');
		TRACE(stream) << TEXT("ShellExecuteEx error ") \
			<< error << TEXT(": ") << buffer << std::endl;
		cerr << stream.str();
	}

	if (info.hProcess != NULL)
	{
		::WaitForSingleObject(info.hProcess, INFINITE);
		::CloseHandle(info.hProcess);
	}
}

#else // !WIN32_API
static std::string& formatPath(std::string& _target, const std::string& _source)
{
	auto position = _source.find(' ', 0);
	decltype(position) offset = 0;

	while (position != _source.npos)
	{
		auto previous = _source.find_last_of('\\', position);
		if (previous == _source.npos) previous = 0;
		else _target.append(_source, offset, ++previous - offset);
		_target.append(1, '\"');

		auto next = _source.find_first_of('\\', position);
		if (next != _source.npos) offset = next + 1;
		else offset = next = _source.size();

		_target.append(_source, previous, next - previous);
		_target.append(1, '\"');
		if (offset != next) _target.append(1, '\\');

		position = _source.find(' ', offset);
	}

	return _target.append(_source, offset, _source.size() - offset);
}

std::string& Structure::getCommand(std::string& _command, IDType _id)
{
	auto count = _counter.fetch_add(1, std::memory_order::relaxed);
	this->count(_id, count);

	auto command = _config._workspace;
	command.append(1, '\\').append(_config._command);

	_command.assign(SHELL_COMMAND);
	formatPath(_command, command).append(" -o\"");
	_command.append(_config._workspace).append(1, '\\');
	_command.append(std::to_string(_id)).append(1, '\"');

	std::string password(_config._length, '0');
	getPassword(password, count);
	return _command.append(" -p").append(password);
}

void Structure::decrypt(IDType _id)
{
	std::string command;
	getCommand(command, _id);

	auto result = std::system(command.c_str());
	if (result != 0)
	{
		std::ostringstream stream;
		stream << '[' << _id << ']';
		TRACEA(stream) << "system error " \
			<< result << std::endl;
		std::cerr << stream.str();
	}
}
#endif // WIN32_API

static auto convert(const std::shared_ptr<void>& _data) noexcept
{
	return std::reinterpret_pointer_cast<Structure>(_data);
}

static auto convert(const std::shared_ptr<void>& _data, \
	std::mutex& _mutex)
{
	std::lock_guard lock(_mutex);
	return std::reinterpret_pointer_cast<Structure>(_data);
}

void Decryptor::move(Decryptor& _left, Decryptor&& _right) noexcept
{
	try
	{
		std::scoped_lock lock(_left._mutex, _right._mutex);
		_left._data = std::move(_right._data);
	}
	catch (std::exception&) {}
}

Decryptor::Decryptor() : _data(nullptr)
{
	static std::once_flag flag;
	std::call_once(flag, initialize);
}

Decryptor& Decryptor::operator=(Decryptor&& _decryptor) noexcept
{
	if (&_decryptor != this)
		move(*this, std::forward<Decryptor>(_decryptor));
	return *this;
}

bool Decryptor::open(const Config& _config)
{
	std::lock_guard lock(_mutex);
	if (_data) return false;

	auto data = std::make_shared<Structure>(_config);
	data->readCount();

	_data = data;
	return true;
}

bool Decryptor::open(Config&& _config)
{
	std::lock_guard lock(_mutex);
	if (_data) return false;

	auto data = std::make_shared<Structure>(std::forward<Config>(_config));
	data->readCount();

	_data = data;
	return true;
}

void Decryptor::close()
{
	std::lock_guard lock(_mutex);
	if (auto data = convert(_data))
	{
		data->writeCount();
		_data.reset();
	}
}

void Decryptor::close(SizeType _count)
{
	std::lock_guard lock(_mutex);
	if (auto data = convert(_data))
	{
		data->writeCount(_count);
		_data.reset();
	}
}

auto Decryptor::count() const -> std::optional<SizeType>
{
	if (auto data = convert(_data, _mutex))
		return data->count();
	return std::nullopt;
}

auto Decryptor::count(IDType _id) const \
-> std::optional<SizeType>
{
	if (auto data = convert(_data, _mutex))
		return data->count(_id);
	return std::nullopt;
}

bool Decryptor::complete() const
{
	auto data = convert(_data, _mutex);
	return data ? data->exist() : true;
}

void Decryptor::execute(IDType _id)
{
	if (auto data = convert(_data, _mutex))
		data->decrypt(_id);
}
